package com.wing.tools.screenlive.encoder;

import android.media.MediaCodec;
import android.media.MediaCodecInfo;
import android.media.MediaFormat;
import android.os.Build;
import android.support.annotation.NonNull;
import android.support.annotation.RequiresApi;
import android.util.Log;
import android.view.Surface;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * Created by wejust on 2018/1/3.
 */

public class VideoEncodeWorker extends Thread {
    public interface IReceiver{
        void onEncodedSample(MediaCodec.BufferInfo info, ByteBuffer data);
    }
    private IReceiver mReceiver;
    //
    private MediaCodec.BufferInfo mBufferInfo;

    // video codec which get access to hardware codec
    private MediaCodec mCodec;

    // indicator for inner loop
    @NonNull
    private final AtomicBoolean mIsRunning = new AtomicBoolean(false);

    private Surface mSurface;

    private final long mTimeoutUsec;

    public VideoEncodeWorker(IReceiver receiver) {
        mReceiver = receiver;

        this.mBufferInfo = new MediaCodec.BufferInfo();
        this.mTimeoutUsec = 10000L;
    }

    public void setIsRunning(boolean running) {
        mIsRunning.set(running);
        if(running){
            mCodec.start();
        }
    }

    public AtomicBoolean isRunning() {
        return mIsRunning;
    }

    public void init(int width, int height, int bitrate, int fps, int frameInterval){
        // configure video output
        MediaFormat format = MediaFormat.createVideoFormat(VideoCodecConstants.VIDEO_CODEC, width, height);
        format.setInteger(MediaFormat.KEY_COLOR_FORMAT,
                MediaCodecInfo.CodecCapabilities.COLOR_FormatSurface);
        format.setInteger(MediaFormat.KEY_BIT_RATE, bitrate);
        format.setInteger(MediaFormat.KEY_FRAME_RATE, fps);
        format.setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, frameInterval);

        try {
            mCodec = MediaCodec.createEncoderByType(VideoCodecConstants.VIDEO_CODEC);
        } catch (IOException e) {
            // can not create avc codec - throw exception
            throw new RuntimeException(e);
        }
        mCodec.configure(format, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE);
        // create surface associated with code
        mSurface = mCodec.createInputSurface();
        //mInputSurface = new CodecInputSurface(mSurface);


        //mSufaceManager = new SurfaceManager(mSurface, TextureEncoder.this.mGLContext, TextureEncoder.this.mGLESVersion);
        // notify codec to start watch surface and encode samples
        //mCodec.start();
    }

    public Surface getSurface(){
        return mSurface;
    }

    @RequiresApi(api = Build.VERSION_CODES.M)
    @Override
    public void run() {
        // prepare video codec
        try {
            encode();
        } finally {
            // release video codec resourses
            release();
        }
    }

    @RequiresApi(api = Build.VERSION_CODES.M)
    void encode() {
        if (!mIsRunning.get()) {
            // if not running anymore, complete stream
            mCodec.signalEndOfInputStream();
        }
        // pre-lollipop api
        if (Build.VERSION.SDK_INT < Build.VERSION_CODES.LOLLIPOP) {

            // get output buffers
            ByteBuffer[] outputBuffers = mCodec.getOutputBuffers();
            for (; ; ) {
                long time = System.currentTimeMillis();
                //get status
                int status = mCodec.dequeueOutputBuffer(mBufferInfo, mTimeoutUsec);
                if (status == MediaCodec.INFO_TRY_AGAIN_LATER) {
                    // something wrong with codec - need try again
                    if (!mIsRunning.get()) {
                        break;
                    }
                } else if (status == MediaCodec.INFO_OUTPUT_BUFFERS_CHANGED) {
                    // need get new output buffers
                    outputBuffers = mCodec.getOutputBuffers();
                } else if (status >= 0) {

                    // encoded sample
                    ByteBuffer data = outputBuffers[status];
                    data.position(mBufferInfo.offset);
                    data.limit(mBufferInfo.offset + mBufferInfo.size);

                    final int endOfStream = mBufferInfo.flags & MediaCodec.BUFFER_FLAG_END_OF_STREAM;
                    if (endOfStream == 0) {
                        mReceiver.onEncodedSample(mBufferInfo, data);
                    }
                    // releasing buffer is important
                    mCodec.releaseOutputBuffer(status, false);

                    // don't have any buffers - need finish
                    if (endOfStream == MediaCodec.BUFFER_FLAG_END_OF_STREAM) {
                        break;
                    }
                }
                Log.d("streaming", "encode time:" + (System.currentTimeMillis() - time));

                try {
                    Thread.sleep(10);
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
        } else {
            for (; ; ) {
                long time = System.currentTimeMillis();
                //get status
                int status = mCodec.dequeueOutputBuffer(mBufferInfo, mTimeoutUsec);
                if (status == MediaCodec.INFO_TRY_AGAIN_LATER) {
                    // something wrong with codec - need try again
                    if (!mIsRunning.get()) {
                        break;
                    }
                } else if (status >= 0) {
                    // encoded sample
                    ByteBuffer data = mCodec.getOutputBuffer(status);
                    if (data != null) {

                        final int endOfStream = mBufferInfo.flags & MediaCodec.BUFFER_FLAG_END_OF_STREAM;
                        if (endOfStream == 0) {
                            mReceiver.onEncodedSample(mBufferInfo, data);
                        }
                        // release buffer
                        mCodec.releaseOutputBuffer(status, false);

                        // don't have any buffers - need finish
                        if (endOfStream == MediaCodec.BUFFER_FLAG_END_OF_STREAM) {
                            break;
                        }
                    }
                }
                //Log.d("streaming", "encode time:" + (System.currentTimeMillis() - time));

                try {
                    Thread.sleep(10);
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
        }
    }

    // release all resources
    public void release() {
        if(mCodec != null){
            mCodec.stop();
            mCodec.release();
        }
    }
}
